#!/usr/bin/env python

from optparse import OptionParser
import os
import os.path
import sys

import charmicat
import histogram

class DrawWorker(charmicat.MPWorker):
    def work(self, hl, opts):        
        for h in hl:
            hname = os.path.join(opts.output, h.name) + opts.ext

            if len(h) == 0:
                with open("{0}_is_empty".format(hname), 'w') as fp:
                    pass

                continue

            if opts.title is None:
                opts.title = h.title

            if len(h.assoc_cols) == 2:
                h.draw(name=hname, fontsize=opts.fontsize,
                       tag=opts.tag, cmap=opts.cmap, title=opts.title)
            else:
                h.draw(name=hname, fontsize=opts.fontsize, 
                       tag=opts.tag, title=opts.title)

        opts.title = None





def main():
    op = OptionParser(usage='%prog [opts] file.pkl')


    op.add_option("-f", "--force", action="store_true", dest="force", 
                  default=False, help="Force overwrite of output directory")
    op.add_option("-l", "--logdir", dest='logdir', default=None,
                  help='Specify directory PATH for MP logging', metavar='PATH')
    op.add_option("-n", "--name", action='append', dest='names', default=[],
                  help='Write images only for histogram NAME (can be repeated)',
                  metavar='NAME')
    op.add_option("-o", "--output-dir", dest="output", default=None, 
                  help="Place output in directory DIR", metavar="DIR")
    op.add_option("-p", "--presentation", action="store_true", 
                  dest="presentation", default=False, 
                  help="Draw images in presentation mode")
    op.add_option("-t", "--tag", dest='tag', default='',
                  help='Append TAG to all histogram titles',
                  metavar='TAG')
    op.add_option("--cmap", dest='cmap', default='hot', 
                  help='Specify colormap CMAP for 2D hists', metavar='CMAP')
    op.add_option("--cpus", dest="ncpu", type="int", default=1,
                  help="Spawn no more than N threads", metavar="N")
    op.add_option("--ext", dest='ext', default='.png', 
                  help='Write out files of type EXT', metavar='EXT')
    op.add_option("--force-log", dest='force_log', default=False, 
                  action='store_true',
                  help='Force reuse of currently-existing log directory')
    op.add_option("--title", dest='title', default=None,
                  help='Set all histogram titles to TITLE', metavar='TITLE')

    opts, args = op.parse_args()

    if len(args) != 1:
        sys.stderr.write("Please supply *one* pickle file with histograms.\n")
        sys.stderr.write("See `draw_hists.py -h' for more.\n")
        sys.stderr.flush()
        return -1

    infname = args[0]

    if opts.output is None:
        opts.output = infname.rsplit('.', 1)[0]

    pngdir = os.path.join(opts.output, 'png')
    pdfdir = os.path.join(opts.output, 'pdf')

    opts.output = charmicat.safe_mkdir(
        opts.output, default=infname.rsplit('.',1)[0], force=opts.force)

    opts.fontsize = 'medium'
    if opts.presentation:
        opts.fontsize = 'xx-large'

    if opts.ext[0] != '.':
        opts.ext = '.' + opts.ext

    hc = histogram.load(infname, names=opts.names)

    if type(hc) in histogram.htypes:
        h = hc
        hc = histogram.HistogramCollection()
        hc.add(h)

    hc_to_draw = histogram.HistogramCollection()

    for h in hc.iter1d():
        hc_to_draw.add(h)
    for h in hc.iter2d():
        hc_to_draw.add(h)

    arglists = charmicat.n_groups(opts.ncpu, list(hc_to_draw))

    monitor = charmicat.CounterMonitor(len(arglists), opts.logdir, interval=2)
    workers = [DrawWorker(monitor.logdir, args=(al, opts)) for al in arglists]

    charmicat.mp_deploy_and_monitor(monitor, workers, ncpu=opts.ncpu)



if __name__ == "__main__":
    try:
        sys.exit(main())
    except charmicat.BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
